import numpy as np
from typing import Optional
from adam import Adam

class OutputProjection:
    def __init__(self, embedding_dim: int, vocab_size: int):
        # Xavier/He initialization: std = sqrt(2 / fan_in)
        std = np.sqrt(2.0 / embedding_dim)
        
        self.w_out = np.random.normal(0.0, std, (embedding_dim, vocab_size))
        self.b_out = np.zeros((1, vocab_size))
        self.optimizer = Adam((embedding_dim, vocab_size))
        self.cached_input: Optional[np.ndarray] = None
    
    def layer_type(self) -> str:
        return "OutputProjection"
    
    def forward(self, input_data: np.ndarray) -> np.ndarray:  # input shape is [sequence_length, embedding_dim]
        self.cached_input = input_data.copy()
        return input_data @ self.w_out + self.b_out  # shape is [sequence_length, vocab_size]
    
    def backward(self, grads: np.ndarray, lr: float) -> np.ndarray:  # grads shape is [sequence_length, vocab_size]
        input_data = self.cached_input
        grad_w_out = input_data.T @ grads
        grad_b_out = np.mean(grads, axis=0, keepdims=True)
        
        grad_input = grads @ self.w_out.T
        
        self.optimizer.step(self.w_out, grad_w_out, lr)
        self.b_out -= lr * grad_b_out
        
        return grad_input